#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <unistd.h>
#include <cerrno>
#include <netinet/in.h>
#include <linux/netlink.h>
#include <linux/netfilter/nf_tables.h>
#include <linux/netfilter/nfnetlink.h>
#include <libmnl/libmnl.h>

#define BUFSIZE 4096

int main(int argc, char *argv[])
{
    if (argc != 2) {
        fprintf(stderr, "Usage: %s <ip_address>\n", argv[0]);
        exit(EXIT_FAILURE);
    }

    const char *ip_address = argv[1];

    struct mnl_socket *nl;
    char buf[BUFSIZE];
    struct nlmsghdr *nlh;
    struct sockaddr_nl sa;
    int ret;

    // 初始化Netlink套接字
    nl = mnl_socket_open(NETLINK_FIREWALL);
    if (nl == NULL) {
        perror("mnl_socket_open");
        exit(EXIT_FAILURE);
    }

    // 初始化Netlink地址结构
    memset(&sa, 0, sizeof(sa));
    sa.nl_family = AF_NETLINK;

    // 将Netlink套接字绑定到Netlink地址
    if (bind(mnl_socket_get_fd(nl), (struct sockaddr *) &sa, sizeof(sa)) < 0) {
        perror("bind");
        exit(EXIT_FAILURE);
    }

    // 创建Netlink消息头
    nlh = mnl_nlmsg_put_header(buf);
    nlh->nlmsg_type = NLMSG_DONE;

    // 添加过滤规则到消息中，禁止指定IP地址的通信
    struct nlattr *nest = mnl_attr_nest_start(nlh, NFNL_SUBSYS_NFTABLES);
    mnl_attr_put_strz(nlh, NFULA_RULE_TABLE, "filter");
    mnl_attr_put_strz(nlh, NFULA_RULE_CHAIN, "INPUT");
    mnl_attr_put_strz(nlh, NFTNL_EXPR, "drop");
    mnl_attr_put_strz(nlh, NFTNL_EXPR_NAME, "accept");
    mnl_attr_put_strz(nlh, NFULA_PAYLOAD, "transport");
    mnl_attr_put_u32(nlh, NFULA_RULE_FAMILY, AF_INET);
    mnl_attr_put_u32(nlh, NFULA_RULE_PROTOCOL, IPPROTO_TCP);
    mnl_attr_put_u32(nlh, NFULA_RULE_FLAGS, NFT_CONTINUE);
    mnl_attr_put_u32(nlh, NFULA_RULE_PAD, 0);

    struct nlattr *rule = mnl_attr_nest_start(nlh, NFULA_RULE_COMPAT);
    mnl_attr_put_strz(nlh, NFTA_COMPAT_DREG, ip_address);
    mnl_attr_nest_end(nlh, rule);

    mnl_attr_nest_end(nlh, nest);

    // 将消息发送到内核
    ret = mnl_socket_sendto(nl, buf, nlh->nlmsg_len);
    if (ret < 0) {
        perror("mnl_socket_sendto");
        exit(EXIT_FAILURE);
    }

    printf("IP address %s has been blocked from network communication.\n", ip_address);

    // 关闭Netlink套接字
    mnl_socket_close(nl);

    return 0;
}
